import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from PIL import Image, ImageTk
import cv2
import numpy as np

class WatermarkRemover:
    def __init__(self, root):
        self.root = root
        self.root.title("图片去水印工具")
        self.root.geometry("800x600")
        self.root.configure(bg='#f0f0f0')
        
        # 主框架
        self.main_frame = tk.Frame(root, padx=20, pady=20, bg='#f0f0f0')
        self.main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 左侧面板（图片预览）
        self.left_frame = tk.Frame(self.main_frame, bg='#f0f0f0')
        self.left_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        
        # 预览标签
        self.preview_label = tk.Label(self.left_frame, text="图片预览", bg='#f0f0f0')
        self.preview_label.pack(pady=5)
        
        # 图片显示区域
        self.canvas = tk.Canvas(self.left_frame, bg='#e0e0e0', width=500, height=400)
        self.canvas.pack(pady=5)
        
        # 右侧面板（控制区）
        self.right_frame = tk.Frame(self.main_frame, bg='#f0f0f0', width=200)
        self.right_frame.pack(side=tk.RIGHT, fill=tk.Y, padx=20)
        
        # 文件选择
        tk.Button(self.right_frame, text="选择图片", command=self.select_image,
                  bg='#007bff', fg='white', width=20).pack(pady=10)
        
        # 水印区域选择说明
        tk.Label(self.right_frame, text="水印区域选择", bg='#f0f0f0').pack(pady=5)
        tk.Label(self.right_frame, text="在图片上拖动选择水印区域", bg='#f0f0f0').pack()
        
        # 处理参数
        tk.Label(self.right_frame, text="处理参数", bg='#f0f0f0').pack(pady=10)
        
        # 半径滑块
        tk.Label(self.right_frame, text="修复半径", bg='#f0f0f0').pack()
        self.radius_scale = tk.Scale(self.right_frame, from_=3, to=20,
                                    orient=tk.HORIZONTAL, length=150)
        self.radius_scale.set(5)
        self.radius_scale.pack()
        
        # 处理按钮
        tk.Button(self.right_frame, text="去除水印", command=self.remove_watermark,
                  bg='#28a745', fg='white', width=20).pack(pady=10)
        tk.Button(self.right_frame, text="保存结果", command=self.save_result,
                  bg='#17a2b8', fg='white', width=20).pack(pady=5)
        
        # 状态栏
        self.status_var = tk.StringVar()
        self.status_var.set("准备就绪")
        tk.Label(self.right_frame, textvariable=self.status_var,
                 bg='#f0f0f0', fg='#333').pack(pady=20)
        
        # 初始化变量
        self.image = None
        self.photo = None
        self.original_image = None
        self.processed_image = None
        self.start_x = None
        self.start_y = None
        self.rect_id = None
        
        # 绑定鼠标事件
        self.canvas.bind('<Button-1>', self.start_selection)
        self.canvas.bind('<B1-Motion>', self.update_selection)
        self.canvas.bind('<ButtonRelease-1>', self.end_selection)
    
    def select_image(self):
        file_path = filedialog.askopenfilename(filetypes=[
            ("图片文件", "*.jpg *.jpeg *.png *.bmp *.gif")
        ])
        if file_path:
            self.original_image = cv2.imread(file_path)
            if self.original_image is None:
                messagebox.showerror("错误", "无法加载图片")
                return
            
            # 调整图片大小以适应显示区域
            display_size = (500, 400)
            self.display_image = self.resize_image(self.original_image, display_size)
            
            # 转换为PhotoImage用于显示
            self.photo = self.cv2_to_photo(self.display_image)
            self.canvas.delete("all")
            self.canvas.create_image(0, 0, anchor=tk.NW, image=self.photo)
            self.status_var.set("已加载图片，请选择水印区域")
    
    def resize_image(self, image, size):
        h, w = image.shape[:2]
        aspect = w / h
        if aspect > size[0] / size[1]:
            new_w = size[0]
            new_h = int(new_w / aspect)
        else:
            new_h = size[1]
            new_w = int(new_h * aspect)
        return cv2.resize(image, (new_w, new_h))
    
    def cv2_to_photo(self, cv2_image):
        rgb_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB)
        pil_image = Image.fromarray(rgb_image)
        return ImageTk.PhotoImage(pil_image)
    
    def start_selection(self, event):
        self.start_x = event.x
        self.start_y = event.y
        if self.rect_id:
            self.canvas.delete(self.rect_id)
        self.rect_id = self.canvas.create_rectangle(
            self.start_x, self.start_y, self.start_x, self.start_y,
            outline='red', width=2
        )
    
    def update_selection(self, event):
        if self.rect_id:
            self.canvas.coords(self.rect_id,
                            self.start_x, self.start_y,
                            event.x, event.y)
    
    def end_selection(self, event):
        if self.rect_id:
            coords = self.canvas.coords(self.rect_id)
            self.selected_area = coords
    
    def remove_watermark(self):
        if not hasattr(self, 'selected_area') or not self.original_image is not None:
            messagebox.showwarning("警告", "请先选择图片和水印区域")
            return
        
        # 获取选择区域的坐标
        x1, y1, x2, y2 = map(int, self.selected_area)
        
        # 计算选择区域在原始图片中的对应位置
        scale_x = self.original_image.shape[1] / self.display_image.shape[1]
        scale_y = self.original_image.shape[0] / self.display_image.shape[0]
        
        orig_x1 = int(x1 * scale_x)
        orig_y1 = int(y1 * scale_y)
        orig_x2 = int(x2 * scale_x)
        orig_y2 = int(y2 * scale_y)
        
        # 创建掩码
        mask = np.zeros(self.original_image.shape[:2], dtype=np.uint8)
        mask[orig_y1:orig_y2, orig_x1:orig_x2] = 255
        
        # 使用Inpainting算法去除水印
        radius = self.radius_scale.get()
        self.processed_image = cv2.inpaint(
            self.original_image, mask, radius, cv2.INPAINT_TELEA
        )
        
        # 显示处理后的图片
        processed_display = self.resize_image(self.processed_image, (500, 400))
        self.photo = self.cv2_to_photo(processed_display)
        self.canvas.delete("all")
        self.canvas.create_image(0, 0, anchor=tk.NW, image=self.photo)
        self.status_var.set("水印去除完成")
    
    def save_result(self):
        if self.processed_image is None:
            messagebox.showwarning("警告", "请先处理图片")
            return
        
        file_path = filedialog.asksaveasfilename(defaultextension=".png",
            filetypes=[("PNG文件", "*.png"), ("JPEG文件", "*.jpg")])
        
        if file_path:
            cv2.imwrite(file_path, self.processed_image)
            self.status_var.set(f"已保存结果到: {file_path}")

if __name__ == "__main__":
    root = tk.Tk()
    app = WatermarkRemover(root)
    root.mainloop()